import itertools
from itertools import filterfalse
import operator

from algorithm.list import inplace_map
from number_theory.gcd import gcd
from number_theory.gcd import gcd_list
from number_theory.gcd import inv_mod
from number_theory.nth_root import int_sqrt
from polynomial.gcd import gcd as gcd_poly
from polynomial.gcd import inv_mod as inv_mod_poly
from polynomial.factor.factor_Zp import factor as factor_Zp
from polynomial.factor.factor_Zp import factor_unique as factor_Zp_unique
from polynomial.factor.lifting_tree import LiftingTree
from polynomial.factor.LLL import LLL
from polynomial.polynomial_Q import PolynomialQ
from polynomial.polynomial_Z import PolynomialZ
from polynomial.polynomial_Zn import PolynomialZn
from prime.prime import Prime

__all__ = ["factor"]


def factor(poly):
    if not poly:
        return [(PolynomialZ(), 1)]
    g = gcd_list(poly.coeff)
    if poly.leading_coeff < 0:
        g = -g
    result = [] if g == 1 else [(PolynomialZ(g), 1)]
    for p, deg in square_free(poly):
        result += ((q, deg) for q in factor_unique(p))
    return sorted(result)


def square_free(poly):
    result = []
    f = poly.primitive()
    h = PolynomialQ()
    for deg in itertools.count(0):
        f1 = gcd_poly(f, f.deriv())
        h1 = f // f1
        g = h // h1
        if deg and g.deg > 0:
            result.append((g, deg))
        f = f1
        h = h1
        if h.deg == 0:
            break
    return result


def select_prime(poly):
    for prime in Prime.generate():
        if poly.leading_coeff % prime:
            p = PolynomialZn(prime)(poly)
            if gcd_poly(p, p.deriv()).deg == 0:
                return prime


def norm_2(vector):
    return int_sqrt(sum(x ** 2 for x in vector)) + 1


def landau_mignotte(poly):
    return (norm_2(poly.coeff) * abs(poly.leading_coeff)) << (poly.deg - 1)


def factor_unique(poly):
    prime = select_prime(poly)
    leading_coeff = poly.leading_coeff
    PolyType = PolynomialZn(prime)
    p_poly = PolyType(poly * inv_mod(leading_coeff, prime))
    fac = factor_Zp_unique(p_poly)
    r = len(fac)
    if r == 1:
        return [poly]

    bound = landau_mignotte(poly) << 1
    tree = LiftingTree(poly, fac)
    mod = prime
    while mod < bound:
        tree.lift()
        mod **= 2
    fac = [PolynomialZ(p) for p in tree.lift_result]
    mod = prime ** ((bound.bit_length() - 1) // (prime - 1).bit_length())
    while mod < bound:
        mod *= prime
    for f in fac:
        inplace_map(lambda x: operator.imod(x, mod), f.coeff)

    return combination(poly, fac, mod)


def combination(poly, fac, mod):
    def check_and_divide(p):
        if sum(map(abs, p.coeff)) > norm << degree:
            return None
        q, remainder = divmod(poly, p)
        return None if remainder else (p, q)

    def pre_check():
        coeff = sum(fac[x].coeff[-2] for x in index)
        coeff = reduce_half(coeff * poly.leading_coeff % mod)
        coeff //= gcd(poly.leading_coeff, coeff)
        return abs(coeff) < norm << degree

    def product():
        p = PolynomialZ(poly.leading_coeff)
        for x in index:
            p *= fac[x]
            inplace_map(lambda x: operator.imod(x, mod), p.coeff)
        inplace_map(reduce_half, p.coeff)
        return p.to_primitive()

    def combination_found(q):
        nonlocal index, fac, norm, poly
        result.append(q[0])
        poly = q[1]
        norm = min(norm, norm_2(poly.coeff))
        length = len(index)
        fac = [fac[x] for x in range(len(fac)) if x not in index]
        if index[0] + length - 1 >= len(fac):
            index = list(range(length + 1))
        else:
            for i in range(1, length):
                index[i] = index[0] + i

    def next_combination():
        nonlocal index
        length = len(index)
        m = length - 1
        while index[m] == len(fac) - length + m:
            m -= 1
        if m == -1:
            index = list(range(length + 1))
        else:
            index[m] += 1
            for i in range(m + 1, length):
                index[i] = index[m] + i - m

    result = []
    half = mod >> 1
    reduce_half = lambda x: x - mod if x > half else x
    norm = norm_2(poly.coeff)
    index = [0]
    while len(index) <= len(fac) >> 1:
        degree = sum(fac[x].deg for x in index)
        q = pre_check() and check_and_divide(product())
        if q:
            combination_found(q)
        else:
            next_combination()
    result.append(PolynomialZ(poly))
    return result


def lll(poly, fac, prime, mod):
    def check_and_divide(vector):
        if not vector[-1] or poly.leading_coeff % vector[-1]:
            return None
        p = PolynomialZ(vector).to_primitive()
        if sum(map(abs, p.coeff)) > norm << p.deg:
            return None
        q, r = divmod(poly, p)
        return None if r else (p, q)

    PolyType = PolynomialZn(prime)
    fac.sort()
    half = mod >> 1
    reduce_half = lambda x: x - mod if x > half else x
    result = []
    norm = norm_2(poly.coeff)

    while fac:
        lattice = LLL(fac[-1], mod)
        for j in range(poly.deg - fac[-1].deg):
            lattice.add_vector()
            lattice.reduction()
            v = lattice.short_vector
            q = check_and_divide(v)
            if q:
                result.append(q[0])
                poly = q[1]
                norm = min(norm, norm_2(poly.coeff))
                g_star = PolyType(q[0])
                fac = list(filter(lambda x: g_star % PolyType(x), fac))
                break
        else:
            fac = []
            result.append(poly)

    return result
